# encoding=utf-8
import argparse
import logging
import os
import shutil
import time
import traceback
import numpy as np
import cv2 as cv
import requests

parser = argparse.ArgumentParser()
parser.add_argument("--remote", type=str, default="http://127.0.0.1:30002/run/optovue/five/opti")
args = parser.parse_args()

logger = logging.getLogger()
logger.setLevel(logging.INFO)

abs_file = os.path.join(os.path.dirname(__file__))
from_path = os.path.join(abs_file, "input_dir")
out_path = os.path.join(abs_file, "output_dir")

temp_path = os.path.join(out_path, "temp_data")
raw_path = os.path.join(out_path, "raw_data")
color_path = os.path.join(out_path, "color_data")

os.makedirs(from_path, exist_ok=True)
os.makedirs(temp_path, exist_ok=True)
os.makedirs(raw_path, exist_ok=True)
os.makedirs(color_path, exist_ok=True)


# format to png
for each_file in os.listdir(from_path):
    if each_file.endswith(".png"):
        shutil.copy(os.path.join(from_path, each_file), os.path.join(temp_path, each_file))
    else:
        t_mat = cv.imread(os.path.join(from_path, each_file))
        name, _ext = os.path.splitext(each_file)
        save_name = os.path.join(temp_path, name + ".png")
        cv.imwrite(save_name, t_mat)

# read data from remote
for each_file in os.listdir(temp_path):
    t1 = time.perf_counter()
    f = {"file": open(os.path.join(temp_path, each_file), 'rb')}
    try:
        res = requests.post(args.remote, files=f)
        code = res.status_code
        if code != 200:
            logging.error(res.content)
        else:
            with open(os.path.join(raw_path, each_file), 'wb') as r_f:
                r_f.write(res.content)
    except Exception as e:
        logging.error(traceback.format_exc())
    t2 = time.perf_counter()
    logging.info("{} cost time, {}s".format(each_file, t2 - t1))


# draw color
color_list = [(0, 0, 255), (255, 0, 0), (0, 255, 0), (192, 255, 0), (0, 255, 156)]
for each_file in os.listdir(raw_path):
    ori_img = cv.imread(os.path.join(raw_path, each_file))
    zero_img = np.zeros_like(ori_img, dtype=np.uint8)
    ori_img = ori_img[:, :, 0]

    for i in range(1, 6):
        zero_img[np.where(ori_img == i)] = color_list[i - 1]

    png_save_path = os.path.join(color_path, each_file)
    cv.imwrite(png_save_path, zero_img)
